"""
麻将训练数据收集器
"""
import os
import json
import numpy as np
from typing import Dict, List, Optional, Tuple, Any
import random
import logging
import time

from sichuanmajiang.engine.game import MahjongGame, Player, AIPlayer
from sichuanmajiang.engine.cards import TileSet
from sichuanmajiang.engine.rules import calculate_score, check_win
from sichuanmajiang.engine.utils import (convert_hand_to_features,
                                          convert_tile_to_index,
                                          calculate_efficiency)


class DataCollector:
    """
    麻将训练数据收集器
    """
    def __init__(self, 
                 save_dir: str = './data',
                 game_config: Optional[Dict] = None):
        """
        Args:
            save_dir: 数据保存目录
            game_config: 游戏配置
        """
        self.save_dir = save_dir
        os.makedirs(save_dir, exist_ok=True)
        
        # 默认游戏配置
        self.game_config = game_config or {
            'max_rounds': 8,
            'max_hand_size': 14,
            'use_ai_players': True,
            'ai_strength': 'medium'  # 'weak', 'medium', 'strong'
        }
        
        # 初始化日志
        self.logger = logging.getLogger('DataCollector')
        
        # 数据收集统计
        self.stats = {
            'games_played': 0,
            'data_points_collected': 0,
            'total_moves': 0,
            'avg_moves_per_game': 0
        }
    
    def collect_game_data(self, 
                          num_games: int,
                          save_interval: int = 100,
                          filename_prefix: str = 'game_data') -> List[Dict]:
        """
        收集游戏数据
        
        Args:
            num_games: 要收集的游戏数量
            save_interval: 保存间隔
            filename_prefix: 文件名前缀
            
        Returns:
            收集到的数据列表
        """
        all_data = []
        start_time = time.time()
        
        self.logger.info(f"开始收集 {num_games} 场游戏的数据...")
        
        for game_idx in range(num_games):
            # 生成游戏配置
            game = self._create_game()
            
            # 收集单场游戏数据
            game_data = self._collect_single_game_data(game)
            all_data.extend(game_data)
            
            # 更新统计信息
            self.stats['games_played'] += 1
            self.stats['data_points_collected'] += len(game_data)
            self.stats['total_moves'] += len(game_data)
            self.stats['avg_moves_per_game'] = self.stats['total_moves'] / self.stats['games_played']
            
            # 日志记录
            if (game_idx + 1) % 10 == 0:
                elapsed = time.time() - start_time
                games_per_minute = (game_idx + 1) / (elapsed / 60)
                self.logger.info(f"已收集 {game_idx + 1}/{num_games} 场游戏, "
                                f"当前数据点: {self.stats['data_points_collected']}, "
                                f"平均每场游戏步数: {self.stats['avg_moves_per_game']:.2f}, "
                                f"速度: {games_per_minute:.2f} 场/分钟")
            
            # 定期保存
            if (game_idx + 1) % save_interval == 0 or (game_idx + 1) == num_games:
                self.save_data(all_data, f"{filename_prefix}_{game_idx + 1}.json")
        
        self.logger.info(f"数据收集完成！总数据点: {self.stats['data_points_collected']}")
        return all_data
    
    def _create_game(self) -> MahjongGame:
        """
        创建一场新的麻将游戏
        """
        # 创建玩家
        players = []
        ai_strength = self.game_config['ai_strength']
        
        for i in range(4):
            if self.game_config['use_ai_players']:
                player = AIPlayer(f"Player_{i}", strategy=ai_strength)
            else:
                player = Player(f"Player_{i}")
            players.append(player)
        
        # 创建游戏实例
        game = MahjongGame(players, max_rounds=self.game_config['max_rounds'])
        return game
    
    def _collect_single_game_data(self, game: MahjongGame) -> List[Dict]:
        """
        收集单场游戏的数据
        
        Args:
            game: 游戏实例
            
        Returns:
            游戏数据列表
        """
        game_data = []
        
        try:
            # 初始化游戏
            game.initialize()
            
            # 游戏循环
            while not game.is_game_over():
                # 收集当前玩家的状态数据
                current_player = game.current_player
                
                # 只有当玩家需要打牌时收集数据
                if current_player.hand.size() >= game.max_hand_size:
                    data_point = self._create_data_point(game, current_player)
                    game_data.append(data_point)
                
                # 执行玩家动作
                game.play_round()
                
                # 避免无限循环
                if len(game_data) > 1000:  # 设置最大步数限制
                    self.logger.warning(f"游戏步数超过限制，提前结束")
                    break
        
        except Exception as e:
            self.logger.error(f"游戏过程中出错: {e}")
            
        return game_data
    
    def _create_data_point(self, game: MahjongGame, player: Player) -> Dict:
        """
        创建一个数据点
        
        Args:
            game: 游戏实例
            player: 当前玩家
            
        Returns:
            数据点字典
        """
        # 提取手牌特征
        hand_features = convert_hand_to_features(player.hand)
        
        # 提取牌河特征
        river_features = np.zeros(34)  # 34种牌
        for p in game.players:
            for tile in p.discarded_tiles:
                tile_idx = convert_tile_to_index(tile)
                river_features[tile_idx] += 1
        
        # 提取其他玩家打出的最后一张牌
        last_discarded = None
        last_discarded_index = -1
        if game.last_discarded_tile:
            last_discarded = game.last_discarded_tile
            last_discarded_index = convert_tile_to_index(last_discarded)
        
        # 计算牌效率
        efficiency = calculate_efficiency(player.hand)
        
        # 计算当前手牌的听牌状态
        tenpai_info = None
        if hasattr(player, 'strategy'):
            tenpai_info = player.strategy.analyze_tenpai(player.hand)
        
        # 创建数据点
        data_point = {
            'timestamp': time.time(),
            'player_id': player.name,
            'round_number': game.current_round,
            'hand_tiles': [str(tile) for tile in player.hand],
            'hand_features': hand_features.tolist(),
            'river_features': river_features.tolist(),
            'last_discarded_tile': str(last_discarded) if last_discarded else None,
            'last_discarded_index': last_discarded_index,
            'remaining_tiles': game.wall.size(),
            'player_wind': player.wind.name,
            'round_wind': game.round_wind.name,
            'score': player.score,
            'efficiency': efficiency,
            'tenpai_info': tenpai_info,
            'valid_actions': self._get_valid_actions(game, player),
            # 动作和结果将在之后填充
            'action': None,
            'reward': None,
            'next_state': None,
            'done': False
        }
        
        # 记录玩家选择的动作
        if hasattr(player, 'next_action') and player.next_action:
            data_point['action'] = player.next_action
        
        return data_point
    
    def _get_valid_actions(self, game: MahjongGame, player: Player) -> List[str]:
        """
        获取当前玩家的有效动作列表
        
        Args:
            game: 游戏实例
            player: 当前玩家
            
        Returns:
            有效动作列表
        """
        valid_actions = []
        
        # 基本打牌动作
        if player.hand.size() >= game.max_hand_size:
            valid_actions.extend(['discard_' + str(tile) for tile in player.hand])
        
        # 如果有最近打出的牌，可以考虑吃碰杠胡
        if game.last_discarded_tile and game.last_player != player:
            # 检查是否可以胡
            if check_win(player.hand, game.last_discarded_tile):
                valid_actions.append('win')
            
            # 检查是否可以杠
            if game.last_discarded_tile in player.hand:
                count = sum(1 for t in player.hand if t == game.last_discarded_tile)
                if count >= 3:
                    valid_actions.append('kong_' + str(game.last_discarded_tile))
            
            # 检查是否可以碰
            if game.last_discarded_tile in player.hand:
                count = sum(1 for t in player.hand if t == game.last_discarded_tile)
                if count >= 2:
                    valid_actions.append('pong_' + str(game.last_discarded_tile))
        
        return valid_actions
    
    def save_data(self, data: List[Dict], filename: str) -> None:
        """
        保存数据到文件
        
        Args:
            data: 要保存的数据
            filename: 文件名
        """
        filepath = os.path.join(self.save_dir, filename)
        
        with open(filepath, 'w', encoding='utf-8') as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
        
        self.logger.info(f"数据已保存到 {filepath}，共 {len(data)} 条记录")
    
    def load_data(self, filename: str) -> List[Dict]:
        """
        从文件加载数据
        
        Args:
            filename: 文件名
            
        Returns:
            加载的数据
        """
        filepath = os.path.join(self.save_dir, filename)
        
        if not os.path.exists(filepath):
            self.logger.error(f"文件不存在: {filepath}")
            return []
        
        with open(filepath, 'r', encoding='utf-8') as f:
            data = json.load(f)
        
        self.logger.info(f"从 {filepath} 加载了 {len(data)} 条记录")
        return data


class ExpertDataCollector(DataCollector):
    """
    专家数据收集器 - 使用模拟的专家策略生成高质量训练数据
    """
    def __init__(self, 
                 save_dir: str = './expert_data',
                 game_config: Optional[Dict] = None,
                 expert_strategy: str = 'aggressive'):  # 'aggressive', 'defensive', 'balanced'
        super().__init__(save_dir, game_config)
        self.expert_strategy = expert_strategy
    
    def _create_game(self) -> MahjongGame:
        """
        创建使用专家策略的游戏
        """
        # 强制使用AI玩家并设置为最强策略
        config = self.game_config.copy()
        config['use_ai_players'] = True
        config['ai_strength'] = 'strong'
        self.game_config = config
        
        return super()._create_game()
    
    def _create_data_point(self, game: MahjongGame, player: Player) -> Dict:
        """
        创建专家数据点，添加更多特征
        """
        data_point = super()._create_data_point(game, player)
        
        # 添加专家级特征
        
        # 计算每张牌的危险度
        danger_scores = self._calculate_danger_scores(game, player)
        data_point['danger_scores'] = danger_scores
        
        # 计算每张牌的期待值
        expected_values = self._calculate_expected_values(player.hand)
        data_point['expected_values'] = expected_values
        
        # 计算防守优先级
        defensive_priority = self._calculate_defensive_priority(game, player)
        data_point['defensive_priority'] = defensive_priority
        
        # 计算进攻优先级
        offensive_priority = self._calculate_offensive_priority(player.hand)
        data_point['offensive_priority'] = offensive_priority
        
        return data_point
    
    def _calculate_danger_scores(self, game: MahjongGame, player: Player) -> List[float]:
        """
        计算每张牌的危险度
        
        Args:
            game: 游戏实例
            player: 当前玩家
            
        Returns:
            各牌的危险度列表
        """
        danger_scores = [0.0] * 34
        
        # 基于其他玩家的弃牌情况计算危险度
        for i, p in enumerate(game.players):
            if p == player:
                continue
            
            # 分析其他玩家的牌河
            for tile in p.discarded_tiles:
                tile_idx = convert_tile_to_index(tile)
                # 根据对手的打牌模式调整危险度
                danger_scores[tile_idx] += 0.1
        
        # 归一化危险度
        max_danger = max(danger_scores) if max(danger_scores) > 0 else 1
        danger_scores = [score / max_danger for score in danger_scores]
        
        return danger_scores
    
    def _calculate_expected_values(self, hand: TileSet) -> List[float]:
        """
        计算每张牌的期待值
        
        Args:
            hand: 当前手牌
            
        Returns:
            各牌的期待值列表
        """
        expected_values = [0.0] * 34
        
        # 基于手牌组合计算每张牌的价值
        # 这里使用简化的计算方法
        for tile in hand:
            tile_idx = convert_tile_to_index(tile)
            # 计算该牌在手中的数量
            count = sum(1 for t in hand if t == tile)
            
            # 孤张价值较低
            if count == 1:
                expected_values[tile_idx] = 0.3
            # 对子价值较高
            elif count == 2:
                expected_values[tile_idx] = 0.7
            # 刻子价值最高
            else:
                expected_values[tile_idx] = 0.9
        
        return expected_values
    
    def _calculate_defensive_priority(self, game: MahjongGame, player: Player) -> float:
        """
        计算防守优先级
        
        Args:
            game: 游戏实例
            player: 当前玩家
            
        Returns:
            防守优先级（0-1之间）
        """
        # 基于对手分数和自己的手牌状态计算防守优先级
        max_opponent_score = max(p.score for p in game.players if p != player)
        own_score = player.score
        
        # 分数差距越大，防守优先级越高
        score_diff = max_opponent_score - own_score
        score_factor = min(score_diff / 1000, 1.0)
        
        # 手牌越差，防守优先级越高
        hand_efficiency = calculate_efficiency(player.hand)
        efficiency_factor = 1.0 - min(hand_efficiency, 1.0)
        
        # 综合计算
        defensive_priority = 0.5 * score_factor + 0.5 * efficiency_factor
        
        # 根据专家策略调整
        if self.expert_strategy == 'defensive':
            defensive_priority = min(defensive_priority * 1.5, 1.0)
        elif self.expert_strategy == 'aggressive':
            defensive_priority = max(defensive_priority * 0.5, 0.0)
        
        return defensive_priority
    
    def _calculate_offensive_priority(self, hand: TileSet) -> float:
        """
        计算进攻优先级
        
        Args:
            hand: 当前手牌
            
        Returns:
            进攻优先级（0-1之间）
        """
        # 基于手牌效率计算进攻优先级
        efficiency = calculate_efficiency(hand)
        offensive_priority = min(efficiency, 1.0)
        
        # 根据专家策略调整
        if self.expert_strategy == 'aggressive':
            offensive_priority = min(offensive_priority * 1.5, 1.0)
        elif self.expert_strategy == 'defensive':
            offensive_priority = max(offensive_priority * 0.5, 0.0)
        
        return offensive_priority


class ExperienceCollector:
    """
    经验收集器 - 用于强化学习中的经验回放
    """
    def __init__(self, 
                 capacity: int = 10000,
                 save_dir: str = './experience'):
        """
        Args:
            capacity: 经验回放池容量
            save_dir: 保存目录
        """
        self.capacity = capacity
        self.save_dir = save_dir
        os.makedirs(save_dir, exist_ok=True)
        
        # 经验回放池
        self.memory = []
        self.position = 0
        
        # 初始化日志
        self.logger = logging.getLogger('ExperienceCollector')
    
    def add_experience(self, 
                       state: np.ndarray,
                       action: int,
                       reward: float,
                       next_state: Optional[np.ndarray] = None,
                       done: bool = False,
                       info: Optional[Dict] = None) -> None:
        """
        添加一条经验
        
        Args:
            state: 当前状态
            action: 执行的动作
            reward: 获得的奖励
            next_state: 下一状态
            done: 是否结束
            info: 额外信息
        """
        experience = {
            'state': state,
            'action': action,
            'reward': reward,
            'next_state': next_state,
            'done': done,
            'info': info or {}
        }
        
        # 如果回放池未满，添加新经验；否则覆盖旧经验
        if len(self.memory) < self.capacity:
            self.memory.append(experience)
        else:
            self.memory[self.position] = experience
        
        # 更新位置指针
        self.position = (self.position + 1) % self.capacity
    
    def sample(self, batch_size: int) -> List[Dict]:
        """
        从经验回放池中随机采样一批经验
        
        Args:
            batch_size: 采样数量
            
        Returns:
            采样的经验列表
        """
        if len(self.memory) < batch_size:
            return self.memory.copy()
        
        return random.sample(self.memory, batch_size)
    
    def save_experience(self, filename: str = 'experience.json') -> None:
        """
        保存经验回放池
        
        Args:
            filename: 文件名
        """
        filepath = os.path.join(self.save_dir, filename)
        
        # 将numpy数组转换为列表以便保存
        serializable_memory = []
        for exp in self.memory:
            serializable_exp = exp.copy()
            if isinstance(serializable_exp['state'], np.ndarray):
                serializable_exp['state'] = serializable_exp['state'].tolist()
            if isinstance(serializable_exp['next_state'], np.ndarray):
                serializable_exp['next_state'] = serializable_exp['next_state'].tolist()
            serializable_memory.append(serializable_exp)
        
        with open(filepath, 'w', encoding='utf-8') as f:
            json.dump(serializable_memory, f, ensure_ascii=False, indent=2)
        
        self.logger.info(f"经验已保存到 {filepath}，共 {len(self.memory)} 条记录")
    
    def load_experience(self, filename: str = 'experience.json') -> None:
        """
        加载经验回放池
        
        Args:
            filename: 文件名
        """
        filepath = os.path.join(self.save_dir, filename)
        
        if not os.path.exists(filepath):
            self.logger.error(f"文件不存在: {filepath}")
            return
        
        with open(filepath, 'r', encoding='utf-8') as f:
            serializable_memory = json.load(f)
        
        # 将列表转换回numpy数组
        self.memory = []
        for exp in serializable_memory:
            memory_exp = exp.copy()
            memory_exp['state'] = np.array(memory_exp['state'])
            if memory_exp['next_state'] is not None:
                memory_exp['next_state'] = np.array(memory_exp['next_state'])
            self.memory.append(memory_exp)
        
        self.position = len(self.memory) % self.capacity
        self.logger.info(f"从 {filepath} 加载了 {len(self.memory)} 条经验记录")
    
    def clear(self) -> None:
        """
        清空经验回放池
        """
        self.memory = []
        self.position = 0
        self.logger.info("经验回放池已清空")
    
    def __len__(self) -> int:
        """
        返回经验回放池中的经验数量
        """
        return len(self.memory)